# 导入必要的模块
import time
import argparse
import os
from photo_processor import IDPhotoMaker
from photo_processor import BackgroundRemovalMethod
from datetime import datetime

def process_photo(input_path, output_path=None, bg_color='blue', size='1inch', width=None, height=None, method='rembg'):
    """
    处理照片并返回输出路径和处理时间
    """
    start_time = time.time()  # 开始计时
    
    # 根据选择的方法创建实例
    photo_maker = IDPhotoMaker(BackgroundRemovalMethod(method))
    
    # 如果提供了自定义尺寸，则使用自定义尺寸
    if width and height:
        photo_maker.ID_PHOTO_WIDTH = width
        photo_maker.ID_PHOTO_HEIGHT = height
    else:
        # 使用预设尺寸
        photo_maker.ID_PHOTO_WIDTH = photo_maker.PHOTO_SIZES[size][0]
        photo_maker.ID_PHOTO_HEIGHT = photo_maker.PHOTO_SIZES[size][1]
    
    # 处理输出路径
    if output_path is None:
        # 使用默认命名
        output_path = _generate_default_filename(input_path, method, bg_color, size, width, height)
    else:
        # 处理输出路径
        output_path = _process_output_path(output_path, input_path, method, bg_color, size, width, height)
    
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    # 处理照片并获取时间统计
    timing = photo_maker.process_photo(input_path, output_path=output_path, bg_color=bg_color)
    
    # 计算总时间
    total_time = time.time() - start_time
    timing['总时间'] = total_time
    
    return output_path, timing

def _process_output_path(output_path, input_path, method, bg_color, size, width, height):
    """处理输出路径，支持全路径、相对路径和目录"""
    # 展开用户路径（如 ~/Documents）
    output_path = os.path.expanduser(output_path)
    
    # 转换为绝对路径
    output_path = os.path.abspath(output_path)
    
    # 如果是目录，在目录下生成默认文件名
    if os.path.isdir(output_path) or output_path.endswith('/') or output_path.endswith('\\'):
        default_name = _generate_default_filename(input_path, method, bg_color, size, width, height)
        return os.path.join(output_path, os.path.basename(default_name))
    
    # 如果没有扩展名，添加 .jpeg 扩展名
    if not os.path.splitext(output_path)[1]:
        output_path += '.jpeg'
    
    return output_path

def _generate_default_filename(input_path, method, bg_color, size, width, height):
    """生成默认的输出文件名"""
    # 获取输入文件名（不含扩展名）
    input_basename = os.path.splitext(os.path.basename(input_path))[0]
    
    # 生成时间戳
    timestamp = time.strftime('%Y%m%d%H%M%S')
    
    # 构建尺寸信息
    size_info = size
    if width and height:
        size_info = f"{width}x{height}"
    
    # 构建输出文件名
    filename = f"output_{input_basename}_{method}_{bg_color}_{size_info}_{timestamp}.jpeg"
    
    # 返回当前目录下的路径
    return os.path.join(os.getcwd(), filename)

def process_directory(input_dir, output_dir=None, **kwargs):
    """处理整个目录中的图片"""
    # 支持的图片格式
    IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
    
    # 如果未指定输出目录，创建默认输出目录
    if output_dir is None:
        output_dir = input_dir.rstrip('/\\') + '_outpath'
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取所有图片文件
    image_files = []
    for root, _, files in os.walk(input_dir):
        for file in files:
            if os.path.splitext(file.lower())[1] in IMAGE_EXTENSIONS:
                image_files.append(os.path.join(root, file))
    
    if not image_files:
        print(f"警告: 在目录 '{input_dir}' 中未找到支持的图片文件")
        return
    
    # 处理进度统计
    total_files = len(image_files)
    processed_files = 0
    failed_files = []
    total_time = 0
    
    print(f"\n开始处理目录: {input_dir}")
    print(f"找到 {total_files} 个图片文件")
    
    # 处理每个图片
    for image_file in image_files:
        processed_files += 1
        rel_path = os.path.relpath(image_file, input_dir)
        print(f"\n[{processed_files}/{total_files}] 处理: {rel_path}")
        
        try:
            # 保持相对路径结构
            rel_dir = os.path.dirname(rel_path)
            output_subdir = os.path.join(output_dir, rel_dir)
            os.makedirs(output_subdir, exist_ok=True)
            
            # 处理单个图片
            output_path, timing = process_photo(
                input_path=image_file,
                output_path=os.path.join(output_subdir, os.path.basename(image_file)),
                **kwargs
            )
            
            total_time += timing['总时间']
            print(f"完成: {output_path}")
            
        except Exception as e:
            print(f"失败: {str(e)}")
            failed_files.append(rel_path)
    
    # 打印处理总结
    print("\n处理完成!")
    print(f"总共处理: {total_files} 个文件")
    print(f"成功: {total_files - len(failed_files)} 个")
    print(f"失败: {len(failed_files)} 个")
    print(f"总耗时: {total_time:.2f} 秒")
    
    if failed_files:
        print("\n处理失败的文件:")
        for file in failed_files:
            print(f"- {file}")

def main():
    parser = argparse.ArgumentParser(
        description='''智能证件照处理工具
        
这个工具可以自动处理证件照，包括去除背景、更换背景颜色、调整尺寸等功能。
支持处理单个图片或整个目录。
        ''',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    # 修改输入路径描述
    parser.add_argument(
        'input_path',
        help='输入图片或目录的路径'
    )
    
    parser.add_argument(
        '--color', '-c',
        default='blue',
        choices=['red', 'blue', 'white', 'gradient'],
        help='''背景颜色选项:
        red - 标准红色背景
        blue - 标准蓝色背景
        white - 纯白色背景
        gradient - 渐变背景
        (默认: blue)'''
    )
    
    parser.add_argument(
        '--size', '-s',
        default='1inch',
        choices=['1inch', '2inch', 'small1inch', 'small2inch'],
        help='''照片尺寸规格:
        1inch - 一寸照片 (25×35mm, 295×413px)
        2inch - 二寸照片 (35×49mm, 413×579px)
        small1inch - 小一寸 (22×32mm, 260×378px)
        small2inch - 小二寸 (35×45mm, 413×531px)
        (默认: 1inch)'''
    )
    
    parser.add_argument(
        '--width', '-w',
        type=int,
        help='自定义输出图片宽度（像素），需要与 --height 一起使用'
    )
    
    parser.add_argument(
        '--height', '-ht',
        type=int,
        help='自定义输出图片高度（像素），需要与 --width 一起使用'
    )

    parser.add_argument(
        '--method', '-m',
        default='rembg',
        choices=['rembg', 'mediapipe', 'deeplab'],
        help='背景去除方法: rembg/mediapipe/deeplab (默认: rembg)'
    )

    # 添加输出路径参数
    parser.add_argument(
        '--output', '-o',
        help='输出文件路径 (默认: 自动生成)'
    )

    # 添加示例用法
    parser.epilog = '''
示例:
  基本用法:
    python example.py input.jpg
  
  指定背景颜色:
    python example.py input.jpg --color blue
    python example.py input.jpg -c white
  
  指定照片尺寸:
    python example.py input.jpg --size 2inch
    python example.py input.jpg -s small1inch
  
  自定义尺寸:
    python example.py input.jpg --width 413 --height 579
    
注意事项:
  1. 输入图片建议使用正面免冠照片
  2. 照片背景最好是纯色或简单背景
  3. 确保人像居中且光线充足
'''
    
    args = parser.parse_args()
    
    # 检查输入路径是否存在
    if not os.path.exists(args.input_path):
        print(f"错误: 找不到输入路径 '{args.input_path}'")
        return
    
    # 提取通用参数
    process_args = {
        'bg_color': args.color,
        'size': args.size,
        'width': args.width,
        'height': args.height,
        'method': args.method
    }
    
    try:
        if os.path.isdir(args.input_path):
            # 处理目录
            process_directory(
                input_dir=args.input_path,
                output_dir=args.output,
                **process_args
            )
        else:
            # 处理单个文件
            output_path, timing = process_photo(
                input_path=args.input_path,
                output_path=args.output,
                **process_args
            )
            print("\n处理完成！")
            print("\n处理时间统计:")
            print(f"{'步骤':<10} {'时间(秒)':<8}")
            print("-" * 20)
            for step, t in timing.items():
                if step != '总时间':
                    print(f"{step:<10} {t:>6.2f}s")
            print("-" * 20)
            print(f"{'总时间':<10} {timing['总时间']:>6.2f}s")
            print(f"\n输出文件: {output_path}")
    except Exception as e:
        print(f"处理失败: {str(e)}")

if __name__ == "__main__":
    main()
